import numpy as np
import bnlearn as bn
import random
import networkx as nx
import matplotlib.pyplot as plt


class Random_Box:
    def __init__(self, row_num, col_num):
        self.row_num = row_num
        self.col_num = col_num
        # 构建row*col种可能，存为一个list
        self.box = set()
        for i in range(row_num):
            for j in range(col_num):
                self.box.add((i, j))

    def add_invalid(self, row, col):
        try:
            self.box.remove((row, col))
        except:
            pass

    def get_random(self):
        if len(self.box) == 0:
            return None
        return random.choice(list(self.box))

    def set_all_invalid(self):
        self.box = set()

    def add_valid(self, row, col):
        self.box.add((row, col))

    def not_empty(self):
        return len(self.box) != 0


class DataConversion:
    ROW_NUM = 10
    NameList = [f"G{i+1}" for i in range(ROW_NUM)]
    D3NameList = (
        [f"A{i+1}" for i in range(ROW_NUM)]
        + [f"B{i+1}" for i in range(ROW_NUM)]
        + [f"C{i+1}" for i in range(ROW_NUM)]
    )

    def general_vector2solutions(vector, mode):
        dim_num = int(mode[1])
        # union
        union = [0] * 100
        for i in range(100):
            for k in range(dim_num):
                if vector[i + k * 100] == 1:
                    union[i] = 1
                    break
        yield ("union", union)

        # intersection
        intersection = [0] * 100
        for i in range(100):
            for k in range(dim_num):
                if vector[i + k * 100] == 0:
                    break
            else:
                intersection[i] = 1
        yield ("intersection", intersection)
        # num -> tail
        for k in range(dim_num):
            cut_solution = vector[k * 100 : (k + 1) * 100]
            yield (f"{chr(65+k)}{chr(65+dim_num-1)}", cut_solution)

    def three_vector2single_vector(vector, policy: str):
        res = [0] * 100
        if policy == "intersection":
            for i in range(100):
                res[i] = (
                    1 if vector[i] == vector[100 + i] == vector[200 + i] == 1 else 0
                )
        elif policy == "union":
            for i in range(100):
                res[i] = (
                    1
                    if vector[i] == 1 or vector[100 + i] == 1 or vector[200 + i] == 1
                    else 0
                )
        elif policy == "AC":
            res = vector[:100]
        elif policy == "BC":
            res = vector[100:200]
        elif policy == "CC":
            res = vector[200:300]
        else:
            print("erorr")
        return res

    def general_vector2matrix(vector, mode):
        if mode == "D1":
            matrix = np.reshape(vector, (10, 10)).astype(int)
        elif mode == "D2":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            res = np.zeros((20, 20))
            res[0:10, 10:20] = matrix1
            res[10:20, 10:20] = matrix2
            matrix = res.astype(int)
        elif mode == "D3":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            matrix3 = DataConversion.vector2matrix(vector[200:300])
            res = np.zeros((30, 30))
            res[0:10, 20:30] = matrix1
            res[10:20, 20:30] = matrix2
            res[20:30, 20:30] = matrix3
            matrix = res.astype(int)
        elif mode == "D4":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            matrix3 = DataConversion.vector2matrix(vector[200:300])
            matrix4 = DataConversion.vector2matrix(vector[300:400])
            res = np.zeros((40, 40))
            res[0:10, 30:40] = matrix1
            res[10:20, 30:40] = matrix2
            res[20:30, 30:40] = matrix3
            res[30:40, 30:40] = matrix4
            matrix = res.astype(int)
        elif mode == "D5":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            matrix3 = DataConversion.vector2matrix(vector[200:300])
            matrix4 = DataConversion.vector2matrix(vector[300:400])
            matrix5 = DataConversion.vector2matrix(vector[400:500])
            res = np.zeros((50, 50))
            res[0:10, 40:50] = matrix1
            res[10:20, 40:50] = matrix2
            res[20:30, 40:50] = matrix3
            res[30:40, 40:50] = matrix4
            res[40:50, 40:50] = matrix5
            matrix = res.astype(int)
        elif mode == "D6":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            matrix3 = DataConversion.vector2matrix(vector[200:300])
            matrix4 = DataConversion.vector2matrix(vector[300:400])
            matrix5 = DataConversion.vector2matrix(vector[400:500])
            matrix6 = DataConversion.vector2matrix(vector[500:600])
            res = np.zeros((60, 60))
            res[0:10, 50:60] = matrix1
            res[10:20, 50:60] = matrix2
            res[20:30, 50:60] = matrix3
            res[30:40, 50:60] = matrix4
            res[40:50, 50:60] = matrix5
            res[50:60, 50:60] = matrix6
            matrix = res.astype(int)
        elif mode == "D7":
            matrix1 = DataConversion.vector2matrix(vector[0:100])
            matrix2 = DataConversion.vector2matrix(vector[100:200])
            matrix3 = DataConversion.vector2matrix(vector[200:300])
            matrix4 = DataConversion.vector2matrix(vector[300:400])
            matrix5 = DataConversion.vector2matrix(vector[400:500])
            matrix6 = DataConversion.vector2matrix(vector[500:600])
            matrix7 = DataConversion.vector2matrix(vector[600:700])
            res = np.zeros((70, 70))
            res[0:10, 60:70] = matrix1
            res[10:20, 60:70] = matrix2
            res[20:30, 60:70] = matrix3
            res[30:40, 60:70] = matrix4
            res[40:50, 60:70] = matrix5
            res[50:60, 60:70] = matrix6
            res[60:70, 60:70] = matrix7
            matrix = res.astype(int)
        return matrix

    def idx2node(idx):
        if idx < 10:
            return f"A{idx+1}"
        elif idx < 20:
            return f"B{idx-9}"
        elif idx < 30:
            return f"C{idx-19}"
        elif idx < 40:
            return f"D{idx-29}"
        elif idx < 50:
            return f"E{idx-39}"
        elif idx < 60:
            return f"F{idx-49}"
        else:
            raise ValueError("idx out of range")

    def general_matrix2edges(matrix, mode):
        edges = []

        for i in range(matrix.shape[0]):
            for j in range(matrix.shape[1]):
                if matrix[i, j] >= 0.5:
                    edges.append(
                        (DataConversion.idx2node(i), DataConversion.idx2node(j))
                    )

        return edges

    def general_vector2edges(vector, mode):
        matrix = DataConversion.general_vector2matrix(vector, mode)
        edges = DataConversion.general_matrix2edges(matrix, mode)
        return edges

    def general_edges2matrix(edges, mode):
        dim_num = int(mode[1])
        matrix = np.zeros((10 * dim_num, 10 * dim_num))
        for edge in edges:
            i = DataConversion.node2idx(edge[0])
            j = DataConversion.node2idx(edge[1])
            matrix[i, j] = 1
        return matrix

    def general_matrix2vector(matrix, mode):
        height = matrix.shape[0]
        width = matrix.shape[1]
        vector = []
        for i in range(height // 10):
            vector += (
                matrix[i * 10 : (i + 1) * 10, width - 10 : width].flatten().tolist()
            )
        return vector

    def general_edges2vector(edges, mode):
        matrix = DataConversion.general_edges2matrix(edges, mode)
        vector = DataConversion.general_matrix2vector(matrix, mode)
        return vector

    def vector2matrix(vector):
        matrix = np.reshape(vector, (10, 10)).astype(int)
        return matrix

    def matrix2edges(matrix):
        edges = []
        for i in range(matrix.shape[0]):
            for j in range(matrix.shape[1]):
                if matrix[i, j] >= 0.5:
                    edges.append(
                        (DataConversion.NameList[i], DataConversion.NameList[j])
                    )
        return edges

    def vector2edges(vector):
        matrix = DataConversion.vector2matrix(vector)
        edges = DataConversion.matrix2edges(matrix)
        return edges

    def matrix2vector(matrix):
        vector = matrix.flatten()
        return vector

    def edges2vector(edges):
        matrix = DataConversion.edges2matrix(edges)
        vector = DataConversion.matrix2vector(matrix)
        return vector

    def node2idx(node):
        if node.startswith("A"):
            return int(node[1:]) - 1
        elif node.startswith("B"):
            return int(node[1:]) - 1 + 10
        elif node.startswith("C"):
            return int(node[1:]) - 1 + 20
        elif node.startswith("D"):
            return int(node[1:]) - 1 + 30
        elif node.startswith("E"):
            return int(node[1:]) - 1 + 40
        elif node.startswith("F"):
            return int(node[1:]) - 1 + 50
        else:
            raise ValueError("node not supported")


class DAG_Utils:
    MAX_REPEAT_TIME = 20

    def _general_matrix_valid(matrix, mode):
        edges = DataConversion.general_matrix2edges(matrix, mode)
        try:
            DGA = bn.make_DAG(edges, methodtype="bayes", verbose=1)
        except:
            return False
        return True

    def general_gen_empty_vector(mode):
        if mode.startswith("D"):
            dim_num = int(mode[1])
            return np.zeros(100 * dim_num)
        else:
            raise ValueError("mode not supported")

    def general_gen_random_vector_by_edges(edges, mode, alpha):
        res_vector = DAG_Utils.general_gen_empty_vector(mode)
        dim_num = int(mode[1])
        for node1, node2 in edges:
            i, j = DataConversion.node2idx(node1), DataConversion.node2idx(node2)
            for k in range(dim_num):
                if random.randint(0, 1) == 1:
                    i, j = j, i
                if random.random() < alpha:
                    res_vector[i * 10 + j + k * 100] = 1
        return res_vector

    def general_gen_random_vector(mode):
        if mode.startswith("D"):
            dim_num = int(mode[1])
            vector = DAG_Utils.general_gen_empty_vector(mode)
            for _ in range(10 * dim_num):
                vector = DAG_Utils.general_add_edge(vector, mode)
            for _ in range(5 * dim_num):
                vector = DAG_Utils.general_remove_edge(vector, mode)
            for _ in range(5 * dim_num):
                vector = DAG_Utils.general_reverse_edge(vector, mode)
            for _ in range(5 * dim_num):
                vector = DAG_Utils.general_add_edge(vector, mode)
            return vector
        else:
            raise ValueError("mode not supported")

    def vector_evaluate(pred_vector, ref_vector):
        pred_vector = np.array(pred_vector).astype(int)
        TP = np.sum((pred_vector == 1) & (ref_vector == 1))
        TN = np.sum((pred_vector == 1) & (ref_vector == 0))

        FP = np.sum((pred_vector == 0) & (ref_vector == 0))
        FP -= 10
        FN = np.sum((pred_vector == 0) & (ref_vector == 1))
        Precision = TP / (TP + FP)
        Recall = TP / (TP + FN)
        if Precision + Recall == 0:
            F1 = 0
        else:
            F1 = 2 * Precision * Recall / (Precision + Recall)
        return {
            "TP": int(TP),
            "TN": int(TN),
            "FP": int(FP),
            "FN": int(FN),
            "Precision": float(Precision),
            "Recall": float(Recall),
            "F1_score": float(F1),
        }

    def general_add_edge(vector, mode):

        repeat_time = 0
        zero_idx = np.where(vector == 0)[0]
        while repeat_time < DAG_Utils.MAX_REPEAT_TIME and len(zero_idx) > 0:
            idx = random.choice(zero_idx)
            vector[idx] = 1
            if DAG_Utils._general_matrix_valid(
                DataConversion.general_vector2matrix(vector, mode), mode
            ):
                break
            else:
                vector[idx] = 0
                zero_idx = np.where(vector == 0)[0]
            repeat_time += 1
        return vector

        dim_num = int(mode[1])
        row_num, col_num = dim_num * 10, dim_num * 10
        random_box = Random_Box(row_num, col_num)
        matrix = DataConversion.general_vector2matrix(vector, mode)

        for i in range(row_num):
            for j in range(col_num):
                if i == j:
                    random_box.add_invalid(i, j)
                elif matrix[i, j] == 1:
                    random_box.add_invalid(i, j)

        repeat_time = 0
        while repeat_time < DAG_Utils.MAX_REPEAT_TIME and random_box.not_empty():
            row, col = random_box.get_random()
            matrix[row, col] = 1
            if DAG_Utils._general_matrix_valid(matrix, mode):
                break
            else:
                matrix[row, col] = 0
                random_box.add_invalid(row, col)
            repeat_time += 1
        return DataConversion.general_matrix2vector(matrix, mode)

    def general_remove_edge(vector, mode):
        one_idx = np.where(vector == 1)[0]
        if len(one_idx) == 0:
            return vector
        idx = random.choice(one_idx)
        vector[idx] = 0
        return vector

        dim_num = int(mode[1])
        row_num, col_num = dim_num * 10, dim_num * 10
        random_box = Random_Box(row_num, col_num)
        matrix = DataConversion.general_vector2matrix(vector, mode)
        random_box.set_all_invalid()
        for i in range(row_num):
            for j in range(col_num):
                if matrix[i, j] == 1:
                    random_box.add_valid(i, j)
        if not random_box.not_empty():
            return vector
        row, col = random_box.get_random()
        matrix[row, col] = 0
        return DataConversion.general_matrix2vector(matrix, mode)

    def general_reverse_edge(vector, mode):
        repeat_time = 0
        one_idx = np.where(vector[200:300] == 1)[0]
        if len(one_idx) == 0:
            return vector
        while repeat_time < DAG_Utils.MAX_REPEAT_TIME and len(one_idx) > 0:
            idx = random.choice(one_idx) + 200
            vector[idx] = 0
            row, col = (idx - 200) // 10, (idx - 200) % 10
            new_edge_idx = col * 10 + row + 200
            vector[new_edge_idx] = 1
            if DAG_Utils._general_matrix_valid(
                DataConversion.general_vector2matrix(vector, mode), mode
            ):
                break
            else:
                vector[idx] = 1
                vector[new_edge_idx] = 0
        return vector

        dim_num = int(mode[1])
        row_num, col_num = dim_num * 10, dim_num * 10
        random_box = Random_Box(row_num, col_num)
        matrix = DataConversion.general_vector2matrix(vector, mode)
        random_box.set_all_invalid()
        for i in range(row_num):
            for j in range(col_num):
                if matrix[i, j] == 1:
                    random_box.add_valid(i, j)
        repeat_time = 0
        while repeat_time < DAG_Utils.MAX_REPEAT_TIME and random_box.not_empty():
            row, col = random_box.get_random()
            matrix[col, row] = 1
            matrix[row, col] = 0
            if DAG_Utils._general_matrix_valid(matrix, mode):
                break
            else:
                matrix[col, row] = 0
                matrix[row, col] = 1
                random_box.add_invalid(row, col)
            repeat_time += 1
        return DataConversion.general_matrix2vector(matrix, mode)

    def general_build_result_vector(vector, mode):
        result_vector = np.zeros(100)
        dim_num = int(mode[1])
        for i in range(100):
            for j in range(dim_num):
                if vector[i + j * 100] == 1:
                    result_vector[i] = 1
                    break
        return result_vector


class DG_Utils:
    def general_show(vector, save_path, mode):
        in_one_result = DAG_Utils.general_build_result_vector(vector, mode)
        edges = DataConversion.vector2edges(in_one_result)
        G = nx.DiGraph()
        G.add_edges_from(edges)
        plt.figure()
        plt.title("Directed Graph of Final Result")
        nx.draw(
            G,
            cmap=plt.get_cmap("jet"),
            with_labels=True,
            node_color="green",
            node_size=400,
        )

        plt.savefig(save_path)
        plt.close()

    def general_mid_show(vector, save_path, mode):
        edges = DataConversion.general_vector2edges(vector, mode)
        DGA = bn.make_DAG(edges, methodtype="bayes", verbose=0)
        out = bn.plot(DGA, verbose=0, title="DAG of Train Result")
        fig = out["fig"]
        fig.savefig(save_path)
        plt.close()


if __name__ == "__main__":
    print("Done~!")
